import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import make_interp_spline, BSpline


def derivative(f, x, h=1e-6):
    return (f(x + h) - f(x - h)) / (2 * h)


def dderivative(f, x, h=1e-6):
    return (derivative(f, x + h, h) - derivative(f, x - h, h)) / (2 * h)


def curvature(f, t):
    return np.linalg.norm(np.cross(derivative(f, t), dderivative(f, t)), 2) / (
        np.linalg.norm(derivative(f, t), 2) ** 3
    )


def interpolate_curve_2(f, x):
    A = np.zeros([3, 3])
    b = np.zeros(3)
    for i in range(3):
        A[i, 0] = 1
        A[i, 1] = x[i]
        A[i, 2] = x[i] ** 2
        b[i] = f(x[i])
    return np.linalg.solve(A, b)


def interpolate_curve3d_2(f, x):
    cx = interpolate_curve_2(lambda t: f(t)[0], x)
    cy = interpolate_curve_2(lambda t: f(t)[1], x)
    cz = interpolate_curve_2(lambda t: f(t)[2], x)
    return np.array([cx, cy, cz])


def gauss_legendre_6point(f, a, b, n):
    nodes, weights = np.polynomial.legendre.leggauss(n)
    x_transformed = 0.5 * (b - a) * nodes + 0.5 * (b + a)
    integral = 0
    for i in range(n):
        integral += weights[i] * f(x_transformed[i])
    return 0.5 * (b - a) * integral


def lambda_func(f, a, b):
    Lt = lambda t: gauss_legendre_6point(
        lambda x: np.linalg.norm(derivative(f, x), 2), a, t, 64
    )
    Kt = lambda t: gauss_legendre_6point(lambda x: curvature(f, x), a, t, 64)
    return lambda t: Lt(t) * 1 + Kt(t) * 0


def root_lambda_func(f, a, b, t0, t1):
    G = lambda_func(f, a, b)
    F = lambda t: G(t) - (G(t0) + G(t1)) / 2
    dL = lambda t: np.linalg.norm(derivative(f, t), 2)
    dK = lambda t: curvature(f, t)
    dF = lambda t: dL(t) * 1 + dK(t) * 0

    t = (t0 + t1) / 2
    while abs(F(t)) > 1e-6:
        t -= F(t) / dF(t)
    return t


def find_features(f, a, b, n):
    t = np.linspace(a, b, n)
    features = []
    for i in range(1, n - 1):
        c0 = curvature(f, t[i - 1])
        c1 = curvature(f, t[i])
        c2 = curvature(f, t[i + 1])
        if c1 > c0 and c1 > c2:
            approx = interpolate_curve3d_2(f, (t[i - 1], t[i], t[i + 1]))
            c = approx[:, 2]
            b = approx[:, 1]
            features.append(-0.5 * np.dot(c, b) / np.dot(c, c))

    return features


def auxiliary_points(f, a, b, features, n):
    G = lambda_func(f, a, b)
    while len(features) < n:
        m = len(features)

        max_diff = 0
        max_idx = 0
        for i in range(1, m):
            diff = G(features[i]) - G(features[i - 1])
            if diff > max_diff:
                max_diff = diff
                max_idx = i

        t = root_lambda_func(f, a, b, features[max_idx - 1], features[max_idx])
        features.insert(max_idx, t)

    return features


def uniform_sample(f, a, b, n):
    t = np.linspace(a, b, n)
    return t, np.array([f(t) for t in t])


def feature_sample(f, a, b, presample_num, n):
    features = find_features(f, a, b, presample_num)
    features.sort()
    features.insert(0, a)
    features.append(b)
    features = auxiliary_points(f, a, b, features, n)
    return features, np.array([f(t) for t in features])


def fitting_err(f, a, b, t, x):
    n = len(t)
    spline_feature_x = make_interp_spline(t, x[:, 0], k=3)
    spline_feature_y = make_interp_spline(t, x[:, 1], k=3)
    spline_feature_z = make_interp_spline(t, x[:, 2], k=3)

    t_fine = np.linspace(a, b, 1000)
    err_x = np.linalg.norm(
        np.array([spline_feature_x(t_fine) - f(t_fine)[0] for t_fine in t_fine]), 2
    )

    err_y = np.linalg.norm(
        np.array([spline_feature_y(t_fine) - f(t_fine)[1] for t_fine in t_fine]), 2
    )

    err_z = np.linalg.norm(
        np.array([spline_feature_z(t_fine) - f(t_fine)[2] for t_fine in t_fine]), 2
    )

    return (err_x**2 + err_y**2 + err_z**2) / n, lambda t: np.array(
        [spline_feature_x(t), spline_feature_y(t), spline_feature_z(t)]
    )


def plot_curve(f, a, b, n, label):
    t, x = uniform_sample(f, a, b, n)
    plt.plot(t, x, label=label)
    plt.grid(True)


def plot_curve3d(f, a, b, n, label):
    t, x = uniform_sample(f, a, b, n)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")  # 创建3D坐标轴

    # 绘制曲线，可自定义颜色、线型、标签等
    ax.plot(
        x[:, 0], x[:, 1], x[:, 2], color="r", linestyle="-", linewidth=2, label=label
    )

    # 添加坐标轴标签和标题
    ax.set_xlabel("X Axis")
    ax.set_ylabel("Y Axis")
    ax.set_zlabel("Z Axis")
    ax.set_title("3D Parametric Curve")

    ax.legend()  # 显示图例


def f0(t):
    fx = lambda t: 3 * t**6 + t**5 - 2 * t**4 + 38 * t**3 - 5 * t**2 - 14 * t
    fy = lambda t: t**6 - 12 * t**5 - 2 * t**4 + 2 * t**3 - 7 * t**2 + 13 * t
    fz = lambda t: np.zeros_like(t)
    return np.array([fx(t), fy(t), fz(t)])


def f1(t):
    fx = lambda t: -26 * t**4 + 64 * t**3 - 52 * t**2 + 16 * t - 1
    fy = lambda t: -112 / 3 * t**4 + 224 / 3 * t**3 - 48 * t**2 + 32 / 3 * t
    fz = lambda t: np.zeros_like(t)
    return np.array([fx(t), fy(t), fz(t)])


def f2(t):
    fx = lambda t: np.sin(2 * t) + np.log(5 * t**4 + 2) + 3 * t**2
    fy = lambda t: 3 * np.exp(t**2 - 1) + np.cos(t / 5) + 2 * t**7
    # fz = lambda t: np.sin(t)
    fz = lambda t: np.zeros_like(t)
    return np.array([fx(t), fy(t), fz(t)])


curve = f2
first = -1
last = 1


# f = lambda_func(curve, first, last)
# t = np.linspace(first, last, 1000)
# ft = [f(t_i) for t_i in t]
# plt.plot(t, ft, label="Lambda Function")

# features = find_features(curve, first, last, 40)
# print(features)

t_fine = np.linspace(first, last, 1000)
plot_curve3d(curve, first, last, 1000, "Helix Curve")

t, x = feature_sample(curve, first, last, 10, 10)
for i in range(len(t)):
    plt.plot(x[i, 0], x[i, 1], x[i, 2] + 1, "o", color="g")
plt.plot(x[:, 0], x[:, 1], x[:, 2] + 1, color="g", linestyle="dotted")

e, curve_fit = fitting_err(curve, first, last, t, x)
print(e)

t, x = uniform_sample(curve, first, last, 10)
for i in range(len(t)):
    plt.plot(x[i, 0], x[i, 1], x[i, 2] - 1, "o", color="b")
plt.plot(x[:, 0], x[:, 1], x[:, 2] - 1, color="b", linestyle="dotted")

e, curve_fit = fitting_err(curve, first, last, t, x)
print(e)

plt.legend()

plt.show()
